package ml.features;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

import corpus.SentPolarityCorpus;

import ml.Data;

import parser.ParserAnnotation;
import parser.ParserPipeline;

public class LemmaExtractor extends FeatureExtractor {
	private String[] lemmas;

	public LemmaExtractor(int offset, SentPolarityCorpus corpus,
			ParserPipeline p) {
		super(offset);
		this.lemmas = getBagOfLemmaArray(corpus, p);
	}

	/**
	 * Cached version.
	 */
	public LemmaExtractor(int offset) throws IOException {
		super(offset);
		lemmas = loadList(new File("features/lemma"));
	}

	private Set<String> getBagOfLemma(SentPolarityCorpus corpus,
			ParserPipeline p) {
		TreeSet<String> set = new TreeSet<String>();
		for (String sent : corpus.getCorpus().keySet()) {
			ParserAnnotation pa = corpus.getParsedPA(sent, p);
			for (String lemma : pa.getLemmas(0))
				set.add(lemma);
		}
		return set;
	}

	private String[] getBagOfLemmaArray(SentPolarityCorpus corpus,
			ParserPipeline p) {
		Set<String> set = getBagOfLemma(corpus, p);
		String[] lemma = new String[set.size()];
		return set.toArray(lemma);
	}

	public List<Data> extract(ParserAnnotation pa, int sentence) {
		List<Data> datas = new ArrayList<Data>(pa.getLemmas(sentence)
				.size());
		List<String> sentLemmas = pa.getLemmas(sentence);
		Set<String> searchSet = new TreeSet<String>(sentLemmas);

		for (String w : searchSet) {
			int index = Arrays.binarySearch(lemmas, w);
			if (index >= 0) {
				datas.add(new Data(1.0, offset + index));
			}
		}
		return datas;
	}
		
	@Override
	public List<Data>[] extractSequence(ParserAnnotation pa, int sentence) {
		List<Data>[] datas = new List[nbrColumns()];
		for (int i = 0; i < datas.length; i++)
			datas[i] = new ArrayList();
		List<String> sentLemmas = pa.getLemmas(sentence);

		for (String tokenLemma : sentLemmas) {
			for (int i = 0; i < lemmas.length; i++) {
				String lemma = lemmas[i];
				datas[i].add(new Data(lemma.equals(tokenLemma), i+offset));
			}
		}
		return datas;
	}
	
	public int nbrColumns() {
		return lemmas.length;
	}

	@Override
	public String[] wekaDescription() {
		String[] description = new String[lemmas.length];
		for (int i = 0; i < description.length; i++) {
			description[i] = "\"" + lemmas[i].replace("\"", "\\\\")
					+ "\" {0.0, 1.0}";
		}
		return description;
	}

	@Override
	public void store() throws IOException {
		File out = new File("features/lemma");
		storeSet(Arrays.asList(lemmas), out);
	}
}
